import sys

sys.path.append(".")

from pathlib import Path

import addict
import onnxruntime as ort
from PySide6.QtCore import QObject, Signal
from PySide6.QtWidgets import (
    QApplication,
    QCheckBox,
    QComboBox,
    QDoubleSpinBox,
    QFileDialog,
    QFormLayout,
    QGridLayout,
    QGroupBox,
    QLabel,
    QMainWindow,
    QMessageBox,
    QProgressBar,
    QPushButton,
    QSpinBox,
    QVBoxLayout,
    QWidget,
    QStatusBar,
)

from app.process import DataSource, run_process_image, run_process_video
from src.models.decoder import Decoder

COLORS = [
    (255, 0, 0),
    (0, 255, 0),
    (0, 0, 255),
    (255, 255, 0),
    (255, 0, 255),
    (0, 255, 255),
    (0, 0, 0),
    (127, 127, 127),
]


class SingalPlayer(QObject):
    percent = Signal(int)  # 进度信号
    finish = Signal(bool)  # 结束信号


class Player(QObject):
    """
    运行功能的控制器
    """

    def __init__(self, paths, cfg) -> None:
        super().__init__()
        self.signals = SingalPlayer()
        self.paths = paths
        self.cfg = cfg

    def run(self):
        # 初始化数据源，模型，解码
        data_source = DataSource(self.paths["input"])
        model = ort.InferenceSession(self.paths["model"], providers=["CUDAExecutionProvider"])
        decoder = Decoder(self.cfg["conf_thres"], scale=64)

        if data_source.mode == "video":
            run_process_video(data_source, model, decoder, self.cfg, self.signals)
        elif data_source.mode == "image":
            run_process_image(data_source, model, decoder, self.cfg, self.signals)
        self.signals.finish.emit(True)
        return


class Parameters:
    """
    模型运行分辨率参数
    """

    def __init__(self) -> None:
        self.input_sizes = [(448, 256), (512, 288), (576, 320), (640, 384), (768, 448), (1024, 576)]
        self.input_sizes_str = [str(x) for x in self.input_sizes]


def raise_message_box(message):
    """
    弹窗提醒
    """
    dlg = QMessageBox()
    dlg.setWindowTitle("Error")
    dlg.setText(message)
    dlg.setStandardButtons(QMessageBox.Ok)
    dlg.setIcon(QMessageBox.Warning)
    dlg.exec()
    return


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("航拍建筑垃圾识别")
        self.statusbar = QStatusBar(self)
        self.setStatusBar(self.statusbar)
        self.params = Parameters()
        self.paths = {"output": "./resources/demo"}
        self.widgets = self._init_widgets()
        self.layouts = self._init_layouts()
        self._define_signal_connects()

    def set_runner(self, runner: Player):
        """设置于界面对应的控制器"""
        self.runner = runner
        return

    def _init_widgets(self):
        """初始化所有控件"""
        widgets = addict.Dict()
        widgets.group.inputs = QGroupBox("数据+模型")
        widgets.btn_load_single = QPushButton("单个文件", self)
        widgets.btn_load_multi = QPushButton("文件夹", self)
        widgets.btn_load_model = QPushButton("加载模型", self)
        widgets.btn_output_dir = QPushButton("输出路径", self)
        widgets.text_single = QLabel("无")
        widgets.text_multi = QLabel("无")
        widgets.text_model = QLabel("无")
        widgets.text_output = QLabel(self.paths["output"])

        widgets.group.params = QGroupBox("参数")
        widgets.input_size_1 = QComboBox(self)
        widgets.input_size_1.addItems(self.params.input_sizes_str)
        widgets.input_size_1.setCurrentIndex(1)
        widgets.input_size_1.setStatusTip("检测模型输入分辨率，越高则能识别到越小的目标，但运行速度越慢")
        widgets.input_size_2 = QComboBox(self)
        widgets.input_size_2.addItems(self.params.input_sizes_str)
        widgets.input_size_2.setCurrentIndex(5)
        widgets.input_size_2.setStatusTip("检测模型输入分辨率，需与前一个分辨率不同")
        widgets.save_mode = QComboBox(self)
        widgets.save_mode.addItems(["image", "video"])
        widgets.save_mode.setStatusTip("结果保存的形式")
        widgets.scale = self._create_spinbox(0.2, 1.0, 0.1, 0.5, "double")
        widgets.scale.setStatusTip("显示缩放尺寸")
        widgets.interval = self._create_spinbox(1, 10, 1, 4, "int")
        widgets.interval.setStatusTip("识别算法运行的帧间隔")
        widgets.threshold = self._create_spinbox(0.1, 0.7, 0.01, 0.25, "double")
        widgets.threshold.setStatusTip("目标置信度阈值，预测值大于该阈值则判定为目标")
        widgets.max_age = self._create_spinbox(1, 3, 1, 1, "int")
        widgets.max_age.setStatusTip("目标跟踪中，最大轨迹延续长度，即目标丢失后延续帧数")
        widgets.min_hit = self._create_spinbox(2, 5, 1, 3, "int")
        widgets.min_hit.setStatusTip("目标跟踪中，建立新轨迹的最小累计帧数")
        widgets.iou_thres = self._create_spinbox(0.1, 0.7, 0.05, 0.3, "double")
        widgets.iou_thres.setStatusTip("目标跟踪中，前后帧目标匹配的 IOU 阈值")
        widgets.min_pixel = self._create_spinbox(2, 100, 1, 10, "int")
        widgets.min_pixel.setStatusTip("可识别的最小目标的像素尺寸")

        widgets.show = QCheckBox(self)

        widgets.group.control = QGroupBox("运行")
        widgets.btn_start = QPushButton("开始", self)
        widgets.text_start = QLabel("在打开实时显示的情况下，如需终止运行，请按Q键")
        widgets.progress = QProgressBar(self)
        widgets.progress.setRange(0, 100)
        return widgets

    def _init_layouts(self):
        """初始化界面布局"""

        # 输入选择区域
        sub_layout_0 = QGridLayout()
        sub_layout_0.addWidget(self.widgets.btn_load_single, 0, 0)
        sub_layout_0.addWidget(self.widgets.text_single, 0, 1)
        sub_layout_0.addWidget(self.widgets.btn_load_multi, 1, 0)
        sub_layout_0.addWidget(self.widgets.text_multi, 1, 1)
        sub_layout_0.addWidget(self.widgets.btn_load_model, 2, 0)
        sub_layout_0.addWidget(self.widgets.text_model, 2, 1)
        sub_layout_0.addWidget(self.widgets.btn_output_dir, 3, 0)
        sub_layout_0.addWidget(self.widgets.text_output, 3, 1)
        self.widgets.group.inputs.setLayout(sub_layout_0)

        # 参数区域
        sub_layout_1 = QFormLayout()
        sub_layout_1.addRow("模型分辨率 1:", self.widgets.input_size_1)
        sub_layout_1.addRow("模型分辨率 2:", self.widgets.input_size_2)
        sub_layout_1.addRow("存储模式(输入为视频):", self.widgets.save_mode)
        sub_layout_1.addRow("输出缩放比:", self.widgets.scale)
        sub_layout_1.addRow("运行间隔(帧):", self.widgets.interval)
        sub_layout_1.addRow("置信度阈值:", self.widgets.threshold)
        sub_layout_1.addRow("跟踪长度(帧):", self.widgets.max_age)
        sub_layout_1.addRow("初始轨迹长度(帧):", self.widgets.min_hit)
        sub_layout_1.addRow("IOU阈值:", self.widgets.iou_thres)
        sub_layout_1.addRow("最小目标(像素):", self.widgets.min_pixel)
        sub_layout_1.addRow("实时显示:", self.widgets.show)
        self.widgets.group.params.setLayout(sub_layout_1)

        # 运行区域
        sub_layout_2 = QGridLayout()
        sub_layout_2.addWidget(self.widgets.btn_start, 0, 0)
        sub_layout_2.addWidget(self.widgets.text_start, 0, 1)
        sub_layout_2.addWidget(self.widgets.progress, 1, 0, 1, 3)
        self.widgets.group.control.setLayout(sub_layout_2)

        layout = QVBoxLayout()
        layout.addWidget(self.widgets.group.inputs)
        layout.addWidget(self.widgets.group.params)
        layout.addWidget(self.widgets.group.control)

        cur_widget = QWidget()
        cur_widget.setLayout(layout)
        self.setCentralWidget(cur_widget)
        return layout

    def _create_spinbox(self, minv, maxv, step, inital, boxtype="double"):
        if boxtype == "double":
            box = QDoubleSpinBox(self)
        elif boxtype == "int":
            box = QSpinBox(self)
        else:
            raise NotImplementedError
        box.setMinimum(minv)
        box.setMaximum(maxv)
        box.setSingleStep(step)
        box.setValue(inital)
        return box

    def _load_single_file(self):
        """加载单个文件"""
        file_name = QFileDialog.getOpenFileName(
            self,
            "Load inputs",
            ".",
            "Videos (*.mp4 *.avi *.MP4 *.AVI);;Images (*.png, *.jpg, *.PNG, *.JPG)",
        )
        self.widgets.text_single.setText(file_name[0])
        self.widgets.text_multi.setText("无")
        self.paths["input"] = file_name[0]
        return

    def _load_multi_file(self):
        """加载文件夹"""
        file_name = QFileDialog.getExistingDirectory(
            self, "Load inputs", ".", QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks
        )
        self.widgets.text_multi.setText(file_name)
        self.widgets.text_single.setText("无")
        self.paths["input"] = file_name
        return

    def _load_model_file(self):
        file_name = QFileDialog.getOpenFileName(self, "Load a model", ".", "Onnx files (*.onnx))")
        self.widgets.text_model.setText(file_name[0])
        self.paths["model"] = file_name[0]
        return

    def _set_output_file(self):
        file_name = QFileDialog.getExistingDirectory(
            self, "Output", ".", QFileDialog.ShowDirsOnly | QFileDialog.DontResolveSymlinks
        )
        self.widgets.text_output.setText(file_name)
        self.paths["output"] = file_name
        return

    def _start_process(self):
        """启动运行"""
        idx_size_1 = self.widgets.input_size_1.currentIndex()
        idx_size_2 = self.widgets.input_size_2.currentIndex()
        if idx_size_1 == idx_size_2:
            raise_message_box("模型的两个输入分辨率需不同")
            return

        if "model" not in self.paths:
            raise_message_box("未选择模型权重文件")
            return

        if "input" not in self.paths:
            raise_message_box("未选择输入数据")
            return

        output_path = Path(self.paths["output"])
        output_path.mkdir(parents=True, exist_ok=True)

        cfg = {
            "output": output_path,
            "save_mode": self.widgets.save_mode.currentText(),
            "scale_g": self.widgets.scale.value(),
            "interval": self.widgets.interval.value(),
            "show": self.widgets.show.isChecked(),
            "sizes": [self.params.input_sizes[idx_size_1], self.params.input_sizes[idx_size_2]],
            "conf_thres": self.widgets.threshold.value(),
            "max_age": self.widgets.max_age.value(),
            "min_hit": self.widgets.min_hit.value(),
            "iou_thres": self.widgets.iou_thres.value(),
            "min_pixel": self.widgets.min_pixel.value(),
        }
        self.widgets.btn_start.setEnabled(False)
        self.player = Player(self.paths, cfg)
        self._define_signal_connects_player()
        self.player.run()
        return

    def _define_signal_connects_player(self):
        """设置控制器的信号与槽的连接"""
        self.player.signals.percent.connect(self.widgets.progress.setValue)
        self.player.signals.finish.connect(self.widgets.btn_start.setEnabled)
        return

    def _define_signal_connects(self):
        """设置界面控件的动作"""
        self.widgets.btn_load_single.clicked.connect(self._load_single_file)
        self.widgets.btn_load_multi.clicked.connect(self._load_multi_file)
        self.widgets.btn_load_model.clicked.connect(self._load_model_file)
        self.widgets.btn_output_dir.clicked.connect(self._set_output_file)
        self.widgets.btn_start.clicked.connect(self._start_process)
        return


if __name__ == "__main__":
    app = QApplication()
    window = MainWindow()
    window.show()

    sys.exit(app.exec())
